{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9e089215-da34-4708-89f7-0f2e15a75e10",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f946bf15-e596-4dd9-861d-edbb6bb984f4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "vocab_size = 16000\n",
    "seq_len = 128\n",
    "d_model = 128\n",
    "n_layer = 4\n",
    "n_head = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "21432ea8-6b23-4549-980d-517e7748c8b2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gpt => 郭靖一掌挥出恶毒果是一技谄的话不加理会众姬犒原本赏去甚是畅快油腻一灯大师的两个嗤的一声点点头延倚奈\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "from torchinfo import summary\n",
    "from torch.nn import functional as F\n",
    "\n",
    "class SinusoidPE(nn.Module):\n",
    "    \"\"\" sin/cos position encoding \"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "        pe = torch.zeros(seq_len, d_model)\n",
    "        pos = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)\n",
    "        emb = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
    "        pe[:,0::2] = torch.sin(pos * emb)\n",
    "        pe[:,1::2] = torch.cos(pos * emb)\n",
    "        \n",
    "        # token embedding: B * C * E\n",
    "        # pos embedding: 1 * C * E\n",
    "        pe = pe.unsqueeze(0)\n",
    "        self.register_buffer('sinusoid_pe', pe)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.sinusoid_pe[:, :x.shape[1],:]\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, n_embd, dropout=0.0):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(n_embd, 4 * n_embd),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4 * n_embd, n_embd),\n",
    "            nn.Dropout(dropout),\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "    \n",
    "    \n",
    "class Head(nn.Module):\n",
    "    \"\"\" one head of self-attention \"\"\"\n",
    "\n",
    "    def __init__(self, head_size, dropout=0.0):\n",
    "        super().__init__()\n",
    "        self.key = nn.Linear(d_model, head_size, bias=False)\n",
    "        self.query = nn.Linear(d_model, head_size, bias=False)\n",
    "        self.value = nn.Linear(d_model, head_size, bias=False)\n",
    "        self.register_buffer('mask', torch.tril(torch.ones(seq_len, seq_len)))\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, C, E = x.shape\n",
    "        k = self.key(x)   # (B, C, E)\n",
    "        q = self.query(x)  # (B, C, E)\n",
    "        # compute attention scores (\"affinities\")\n",
    "        wei = q @ k.transpose(-2, -1) * E**-0.5  # (B, C, E) @ (B, E, C) -> (B, C, C)\n",
    "        wei = wei.masked_fill(self.mask[:C,:C] == 0, float('-inf'))  # (B, C, C)\n",
    "        wei = F.softmax(wei, dim=-1)  # (B, C, C)\n",
    "        wei = self.dropout(wei)\n",
    "        # perform the weighted aggregation of the values\n",
    "        v = self.value(x)  # (B, C, E)\n",
    "        out = wei @ v  # (B, C, C) @ (B, C, E) -> (B, C, E)\n",
    "        return out\n",
    "    \n",
    "class SelfAttention(nn.Module):\n",
    "    def __init__(self, num_heads, head_size, dropout=0.0):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
    "        self.proj = nn.Linear(d_model, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
    "        out = self.dropout(self.proj(out))\n",
    "        return out\n",
    "    \n",
    "class MultiAttention(nn.Module):\n",
    "    def __init__(self, dropout=0.0):\n",
    "        super().__init__()\n",
    "        # self.w_q = nn.Linear(d_model, d_model)\n",
    "        # self.w_k = nn.Linear(d_model, d_model)\n",
    "        # self.w_v = nn.Linear(d_model, d_model)\n",
    "        self.attn = nn.Linear(d_model, 3 * d_model)\n",
    "        self.proj = nn.Linear(d_model, d_model)\n",
    "        self.dropout = nn.Dropout(dropout) \n",
    "        self.register_buffer('mask', torch.tril(torch.ones(seq_len, seq_len))\n",
    "                                          .view(1,1, seq_len, seq_len))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        B, C, E = x.shape\n",
    "        q, k, v = self.attn(x).split(d_model, dim=2)\n",
    "        q = q.view(B, C, n_head, E // n_head).transpose(1,2) # (B, C, nh, hs) -> (B, nh, C, hs)\n",
    "        k = k.view(B, C, n_head, E // n_head).transpose(1,2) # (B, C, nh, hs)\n",
    "        v = v.view(B, C, n_head, E // n_head).transpose(1,2) # (B, C, nh, hs)\n",
    "        \n",
    "        # (B, nh, C, hs) * (B, nh, hs, C) -> (B, nh, C, C)\n",
    "        wei = q @ k.transpose(-2, -1) * (k.size(-1))**-0.5 \n",
    "        wei = wei.masked_fill(self.mask[:,:,:C,:C] == 0, float('-inf'))\n",
    "        wei = F.softmax(wei, dim=-1)\n",
    "        wei = self.dropout(wei)\n",
    "        att = wei @ v # (B, nh, C, C) * (B, nh, C, hs) -> (B, nh, C, hs)\n",
    "        att = att.transpose(1,2).contiguous().view(B,C,E) # (B, nh, C, hs) -> (B, C, nh, hs) -> (B, C, E)\n",
    "        \n",
    "        out = self.proj(att)\n",
    "        return out\n",
    "            \n",
    "class Block(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        head_size = d_model // n_head\n",
    "        self.ln1 = nn.LayerNorm(d_model)\n",
    "        # self.attn = SelfAttention(n_head, head_size)\n",
    "        self.attn = MultiAttention()\n",
    "        self.ln2 = nn.LayerNorm(d_model)\n",
    "        self.ffn = FeedForward(d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x + self.attn(self.ln1(x))\n",
    "        x = x + self.ffn(self.ln2(x))\n",
    "        return x\n",
    "\n",
    "class GPTModel(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.tok_embed_table = nn.Embedding(vocab_size, d_model)\n",
    "        self.pos_embed_table = SinusoidPE()\n",
    "        self.decoder_blocks = nn.Sequential(*[Block() for _ in range(n_layer)])\n",
    "        self.ln = nn.LayerNorm(d_model)\n",
    "        self.final_linear = nn.Linear(d_model, vocab_size)\n",
    "        \n",
    "    def forward(self, features, targets=None):\n",
    "        tok_emb = self.tok_embed_table(features)\n",
    "        pos_emb = self.pos_embed_table(tok_emb)\n",
    "        x = tok_emb + pos_emb\n",
    "        x = self.decoder_blocks(x)\n",
    "        out = self.final_linear(self.ln(x))\n",
    "        \n",
    "        if targets is not None:\n",
    "            B, C, V = out.shape\n",
    "            out = out.view(B * C, V)\n",
    "            targets = targets.view(B * C)\n",
    "            loss = F.cross_entropy(out, targets)\n",
    "            return out, loss\n",
    "        else:\n",
    "            return out, None\n",
    "      \n",
    "    @torch.no_grad()\n",
    "    def generate(self, seq, max_new_tokens):\n",
    "        for _ in range(max_new_tokens):\n",
    "            seq = seq[:,-seq_len:] # B, L, E\n",
    "            pred, _ = self(seq)\n",
    "            pred = pred[:,-1,:] # B, C, V -> B, 1, V\n",
    "            probs = F.softmax(pred, dim=-1)\n",
    "            next_token = torch.multinomial(probs, num_samples=1) # [0.1, 0.7, 0.2]\n",
    "            seq = torch.cat((seq, next_token), dim=1)\n",
    "        return seq\n",
    "        \n",
    "model = GPTModel()\n",
    "# summary(model)\n",
    "\n",
    "import sentencepiece as spm\n",
    "import sys\n",
    "\n",
    "model_file = \"bird_shooter.model\"\n",
    "sp = spm.SentencePieceProcessor()\n",
    "if not sp.load(model_file=model_file):\n",
    "    print(\"load tokenizer model failed\")\n",
    "    sys.exit(1)\n",
    "\n",
    "user_input = \"郭靖一掌挥出\"   \n",
    "context = torch.tensor([sp.encode(user_input)], dtype=torch.int32)\n",
    "gpt_output = model.generate(context, max_new_tokens=20)[0].tolist()\n",
    "print(f\"gpt => {sp.decode(gpt_output)}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
